package main

type AvlTree struct {
	Root *TreeNode
}

func NewAvlTree(num int) *AvlTree {
	return &AvlTree{
		Root: NewTreeNode(num),
	}
}

// Height 获取节点的高度，如果节点为nil，高度为-1
func (a *AvlTree) Height(node *TreeNode) int {
	if node == nil {
		return -1
	}
	return node.Height
}

// UpdateHeight 更新节点的高度，节点的高度等于 子节点中的最高高度+1
func (a *AvlTree) UpdateHeight(node *TreeNode) {
	lh := a.Height(node.Left)
	rh := a.Height(node.Right)
	if lh > rh {
		node.Height = lh + 1
	} else {
		node.Height = rh + 1
	}
}

// BalanceFactor 获取节点的平衡因子。平衡因子=左子节点的高度-右子节点的高度。空节点的平衡因子设置为0
// 一颗符合AVL树定义的节点的平衡因子只能是3种情况，-1（左节点为空，右节点不为空），0（左右节点皆不为空），1（左节点不为空，右节点为空）
func (a *AvlTree) BalanceFactor(node *TreeNode) int {
	if node == nil {
		return 0
	}
	return node.Left.Height - node.Right.Height
}

// RightRotate 节点右旋
func (a *AvlTree) RightRotate(node *TreeNode) *TreeNode {
	// node 是要右旋的节点
	// 向右旋转，说明node是左偏树，拿到旋转节点左节点
	child := node.Left
	// 因为node 要向右旋转，那么先把child的右节点保存起来，避免被覆盖掉
	grandChild := child.Right
	child.Right = node
	node.Left = grandChild
	// 实际上受到影响的节点只有child和node，孙子节点没有影响，只对受影响的节点重新计算高度
	a.UpdateHeight(child)
	a.UpdateHeight(node)
	return child
}

// LeftRotate 节点左旋，和右旋一致，倒换一下方向即可
func (a *AvlTree) LeftRotate(node *TreeNode) *TreeNode {
	child := node.Right
	grandChild := child.Left
	child.Left = node
	node.Right = grandChild
	a.UpdateHeight(child)
	a.UpdateHeight(node)
	return child
}

// Rotate 旋转树的节点，使子树恢复平衡
func (a *AvlTree) Rotate(node *TreeNode) *TreeNode {
	bf := a.BalanceFactor(node)
	// 左偏树
	if bf > 1 {
		if a.BalanceFactor(node.Left) >= 0 {
			// 右旋
			return a.RightRotate(node)
		} else {
			// 先左旋再右旋
			node.Left = a.LeftRotate(node.Left)
			return a.RightRotate(node)
		}
	} else if bf < 1 {
		// 右偏树
		if a.BalanceFactor(node.Right) <= 0 {
			// 左旋
			return a.RightRotate(node)
		} else {
			// 先右旋再左旋
			node.Right = a.RightRotate(node.Right)
			return a.LeftRotate(node)
		}
	}

	// 为0的情况，平衡树不需要旋转
	return node
}

// Insert inserts a new node with the given integer value into the AVL tree.
// It updates the root of the tree after insertion to maintain the AVL property.
//
// Parameters:
//
//	num - the integer value to be inserted into the tree.
func (a *AvlTree) Insert(num int) {
	a.Root = a.insertHelper(a.Root, num)
}

func (a *AvlTree) insertHelper(node *TreeNode, num int) *TreeNode {
	if node == nil {
		return NewTreeNode(num)
	}

	if num < node.Val.(int) {
		// 插入到左侧
		node.Left = a.insertHelper(node.Left, num)
	} else if num > node.Val.(int) {
		// 插入到右侧
		node.Right = a.insertHelper(node.Right, num)
	} else {
		// 如果num 和 node 的值相同，那么直接返回当前node节点即可
		return node
	}

	a.UpdateHeight(node)
	a.Rotate(node)

	return node
}

func (a *AvlTree) Delete(num int) {
	a.Root = a.deleteHelper(a.Root, num)
}

func (a *AvlTree) deleteHelper(node *TreeNode, num int) *TreeNode {
	if node == nil {
		return nil
	}

	if num < node.Val.(int) {
		node.Left = a.deleteHelper(node.Left, num)
	} else if num > node.Val.(int) {
		node.Right = a.deleteHelper(node.Right, num)
	} else {
		// 节点的度为0或者1的时候
		if node.Left == nil || node.Right == nil {
			child := node.Left
			if node.Right != nil {
				child = node.Right
			}

			if child == nil {
				// 子节点数量为0，直接删除该节点
				return nil
			} else {
				// 子节点数量为1，直接删除node
				node = child
			}

			// 节点的度为2的时候，则将中序遍历的下个节点删除，并用该节点替换当前节点
		} else {
			tmp := node.Right
			for tmp.Left != nil {
				tmp = tmp.Left
			}
			node.Right = a.deleteHelper(node.Right, tmp.Val.(int))
			node.Val = tmp.Val

		}
	}

	a.UpdateHeight(node)
	a.Rotate(node)

	return node
}
